import json
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image
import numpy as np
import os

def visualize_single_question(json_file_path, question_index=0):
    """
    可视化单个问题
    
    Args:
        json_file_path: JSON文件路径
        question_index: 问题索引
    """
    # 读取JSON文件
    with open(json_file_path, 'r') as f:
        data = json.load(f)
    
    if question_index >= len(data['questions']):
        print("Question index {} out of range".format(question_index))
        return
    
    question_data = data['questions'][question_index]
    
    # 获取数据
    image_filename = question_data['image_filename']
    question = question_data['question']
    chain_questions = question_data['chain_questions']
    chain_answers = question_data['chain_answers']
    final_answer = question_data['answer']
    
    # 构建图像路径
    image_path = "../output/images/{}".format(image_filename)
    
    # 检查图像是否存在
    if not os.path.exists(image_path):
        print("Image not found: {}".format(image_path))
        return
    
    # 创建图形
    fig = plt.figure(figsize=(20, 12))
    
    # 左侧：显示图像
    ax1 = plt.subplot(1, 2, 1)
    img = Image.open(image_path)
    ax1.imshow(img)
    ax1.set_title("Image: {}".format(image_filename), fontsize=14, fontweight='bold')
    ax1.axis('off')
    
    # 在图像上标记最终答案位置
    if len(final_answer) == 3:  # [x, y, depth]
        x, y, depth = final_answer
        # 绘制最终答案的圆圈
        circle = patches.Circle((x, y), 15, fill=False, edgecolor='red', linewidth=3, label='Final Answer')
        ax1.add_patch(circle)
        
        # 添加深度信息
        ax1.text(x + 20, y, "depth: {:.2f}".format(depth), fontsize=10, color='red', 
                bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
    
    # 右侧：显示问题和Chain of Thought
    ax2 = plt.subplot(1, 2, 2)
    ax2.axis('off')
    
    # 设置背景色
    ax2.set_facecolor('#f0f0f0')
    
    # 显示主问题
    y_pos = 0.95
    ax2.text(0.05, y_pos, "Question:", fontsize=16, fontweight='bold', color='red')
    y_pos -= 0.08
    ax2.text(0.05, y_pos, question, fontsize=12, wrap=True, 
            bbox=dict(boxstyle="round,pad=0.5", facecolor="white", alpha=0.9))
    y_pos -= 0.15
    
    # 显示Chain of Thought过程
    ax2.text(0.05, y_pos, "Chain of Thought:", fontsize=16, fontweight='bold', color='blue')
    y_pos -= 0.08
    
    for i, (chain_q, chain_a) in enumerate(zip(chain_questions, chain_answers)):
        # 显示子问题
        ax2.text(0.05, y_pos, "Step {}: {}".format(i+1, chain_q), fontsize=11, fontweight='bold', color='darkblue')
        y_pos -= 0.06
        
        # 显示子答案
        if len(chain_a) == 3:  # [x, y, depth]
            x, y, depth = chain_a
            answer_text = "Answer: ({}, {}, depth: {:.2f})".format(x, y, depth)
        else:
            answer_text = "Answer: {}".format(chain_a)
        
        ax2.text(0.08, y_pos, answer_text, fontsize=10, color='green',
                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgreen", alpha=0.7))
        y_pos -= 0.12
        
        # 在图像上标记当前步骤的物体位置
        if len(chain_a) == 3:  # [x, y, depth]
            x, y, depth = chain_a
            # 绘制带数字的圆圈
            circle = patches.Circle((x, y), 12, fill=True, facecolor='blue', alpha=0.7)
            ax1.add_patch(circle)
            # 添加步骤数字
            # ax1.text(x, y, str(i+1), fontsize=10, color='white', weight='bold', ha='center', va='center')
            ax1.text(x, y, str(i+1), fontsize=10, color=(1, 1, 1, 0.5), weight='bold', ha='center', va='center')
    
    # 显示最终答案
    ax2.text(0.05, y_pos, "Final Answer:", fontsize=16, fontweight='bold', color='red')
    y_pos -= 0.08
    if len(final_answer) == 3:
        x, y, depth = final_answer
        final_answer_text = "({}, {}, depth: {:.2f})".format(x, y, depth)
    else:
        final_answer_text = str(final_answer)
    
    ax2.text(0.05, y_pos, final_answer_text, fontsize=12, color='red',
            bbox=dict(boxstyle="round,pad=0.5", facecolor="lightcoral", alpha=0.8))
    
    # 调整布局
    plt.tight_layout()
    
    # 显示图形
    plt.show()
    
    # 保存图形
    plt.savefig(f"../visualize/photos/{question_index:06d}.png")


def visualize_questions(json_file_path, question_number=1):
    for i in range(question_number):
        visualize_single_question(json_file_path, i)
    

if __name__ == "__main__":
    # 可视化所有问题
    json_file_path = "../output/xyg_test_cot_len4.json"
    
    visualize_questions(json_file_path, 4)
